from loadDataset import Load2022Dataset, LoadMachineData
from torch.utils.data import Dataset
from typing import Tuple
import numpy as np
import random


class Dataset2022ForMultiUtil:
    @staticmethod
    def decimation_by_avg(data, factor):
        """Function for replacing each sequence of previous factor samples with their average"""
        # for example, frame [0, 800]ms -> 17samples (Krusienski et al., 2006)
        # data.shape = [ch, time, trial]
        ratio_dsample = factor
        n_ch, n_frame, n_trial = data.shape
        # print(n_frame)
        decimated_frame = int(np.floor(n_frame / ratio_dsample))
        # print(decimated_frame)
        # memory pre-allocation
        decimated_data = np.zeros((n_ch, decimated_frame, n_trial))
        # print(decimated_data.shape)
        for i in range(n_trial):
            for j in range(decimated_frame):
                cur_data = data[:, :, i]
                decimated_data[:, j, i] = np.mean(
                    cur_data[:, j * ratio_dsample : (j + 1) * ratio_dsample], axis=1
                )
        return decimated_data


class Dataset2022ForMulti(Dataset):
    def __init__(self, channelsNum: Tuple, matPath: str, useAvg: bool = False) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        # self.signalLen = 307
        # (55, 32, 307, 300)
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        self.tarData, self.noTarData = dataset.getDataNdarray()
        self.signalLen = self.tarData.shape[2]
        self.dataTarSize, self.dataNonTargetSize = dataset.getDataSize()
        self.dataTarArraySize, self.dataNonTargetArraySize = dataset.getDataArraySize()
        self.channelsNum = channelsNum
        self.tarDataArrayIndex = 0
        self.nonTarDataArrayIndex = 0
        # self.tarDataIndex = 0
        # self.nonTarDataIndex = 0
        self.personSignalLen = len(channelsNum)
        self.dataSize = self.dataTarSize + self.dataNonTargetSize
        self.data = np.empty((self.dataSize, 2 * self.personSignalLen, self.signalLen))
        self.label = np.empty(self.dataSize)
        count = 0
        if useAvg:
            self.tarData = self.decimation_by_avg(self.tarData, 24)
            self.noTarData = self.decimation_by_avg(self.noTarData, 24)
        while count < self.dataTarSize:
            subjectA = self.tarData[self.tarDataArrayIndex]
            subjectB = self.tarData[
                (self.tarDataArrayIndex + 1) % self.dataTarArraySize
            ]
            if np.size(subjectA, 2) != np.size(subjectB, 2):
                minLen = min(np.size(subjectA, 2), np.size(subjectB, 2))
            else:
                minLen = np.size(subjectA, 2)
            for i in range(minLen):
                dataAList = []
                dataBList = []
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
                for j in self.channelsNum:
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
                self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
                self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                    np.array(dataBList)
                )
                self.label[count] = 1
                count += 1
            self.tarDataArrayIndex += 1
        while count - self.dataTarSize < self.dataNonTargetSize:
            subjectA = self.noTarData[self.nonTarDataArrayIndex]
            subjectB = self.noTarData[
                (self.nonTarDataArrayIndex + 1) % self.dataNonTargetArraySize
            ]
            if np.size(subjectA, 2) != np.size(subjectB, 2):
                minLen = min(np.size(subjectA, 2), np.size(subjectB, 2))
            else:
                minLen = np.size(subjectA, 2)
            for i in range(minLen):
                dataAList = []
                dataBList = []
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
                for j in self.channelsNum:
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
                self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
                self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                    np.array(dataBList)
                )
                self.label[count] = 0
                count += 1
            self.nonTarDataArrayIndex += 1
        if count < self.dataSize:
            self.data = self.data[0:count, :, :]
            self.dataSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


class Dataset2022TarForMulti(Dataset):
    def __init__(self, channelsNum: Tuple, matPath: str, useAvg: bool = False) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        self.tarData, self._ = dataset.getDataNdarray()
        self.signalLen = self.tarData.shape[2]
        self.dataTarSize, self._ = dataset.getDataSize()
        self.dataTarArraySize, self._ = dataset.getDataArraySize()
        self.channelsNum = channelsNum
        self.personSignalLen = len(channelsNum)
        self.tarDataArrayIndex = 0
        # self.tarDataIndex = 0
        self.data = np.empty(
            (self.dataTarSize, 2 * self.personSignalLen, self.signalLen)
        )
        self.label = np.empty(self.dataTarSize)
        count = 0
        if useAvg:
            self.tarData = self.decimation_by_avg(self.tarData, 24)
        while count < self.dataTarSize:
            subjectA = self.tarData[self.tarDataArrayIndex]
            subjectB = self.tarData[
                (self.tarDataArrayIndex + 1) % self.dataTarArraySize
            ]
            if np.size(subjectA, 2) != np.size(subjectB, 2):
                minLen = min(np.size(subjectA, 2), np.size(subjectB, 2))
            else:
                minLen = np.size(subjectA, 2)
            for i in range(minLen):
                dataAList = []
                dataBList = []
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
                for j in self.channelsNum:
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
                self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
                self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                    np.array(dataBList)
                )
                self.label[count] = 1
                count += 1
            self.tarDataArrayIndex += 1
        if count < self.dataTarSize:
            self.data = self.data[0:count, :, :]
            self.dataTarSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataTarSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


class Dataset2022NonTarForMulti(Dataset):
    def __init__(
        self, channelsNum: Tuple[int, int, int], matPath: str, useAvg: bool = False
    ) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        self._, self.noTarData = dataset.getDataNdarray()
        self.signalLen = self.noTarData.shape[2]
        self._, self.dataNonTargetSize = dataset.getDataSize()
        self._, self.dataNonTargetArraySize = dataset.getDataArraySize()
        self.channelsNum = channelsNum
        self.personSignalLen = len(channelsNum)
        self.nonTarDataArrayIndex = 0
        # self.nonTarDataIndex = 0
        self.data = np.empty(
            (self.dataNonTargetSize, 2 * self.personSignalLen, self.signalLen)
        )
        self.label = np.empty(self.dataNonTargetSize)
        count = 0
        if useAvg:
            self.noTarData = self.decimation_by_avg(self.noTarData, 24)
        while count < self.dataNonTargetSize:
            subjectA = self.noTarData[self.nonTarDataArrayIndex]
            subjectB = self.noTarData[
                (self.nonTarDataArrayIndex + 1) % self.dataNonTargetArraySize
            ]
            if np.size(subjectA, 2) != np.size(subjectB, 2):
                minLen = min(np.size(subjectA, 2), np.size(subjectB, 2))
            else:
                minLen = np.size(subjectA, 2)
            for i in range(minLen):
                dataAList = []
                dataBList = []
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
                for j in self.channelsNum:
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
                self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
                self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                    np.array(dataBList)
                )
                self.label[count] = 0
                count += 1
            self.nonTarDataArrayIndex += 1
        if count < self.dataNonTargetSize:
            self.data = self.data[0:count, :, :]
            self.dataNonTargetSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataNonTargetSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


class Dataset2022ForMultiMean(Dataset):
    def __init__(self, channelsNum: Tuple, matPath: str) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        # self.signalLen = 307
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        self.tarData, self.noTarData = dataset.getDataNdarray()
        self.signalLen = self.tarData.shape[2]
        # self.tarDataMean = np.empty((self.tarData.shape[0], self.tarData.shape[1], self.tarData.shape[2]))
        # self.noTarDataMean = np.empty((self.noTarData.shape[0], self.noTarData.shape[1], self.noTarData.shape[2]))
        self.dataTarArraySize = self.tarData.shape[0]
        self.dataNonTargetArraySize = self.noTarData.shape[0]
        # for i in range(self.tarData.shape[0]):
        #     self.tarDataMean[i] = np.mean(self.tarData[i], axis=2)
        # for i in range(self.noTarData.shape[0]):
        #     self.noTarDataMean[i] = np.mean(self.noTarData[i], axis=2)
        self.tarDataMean = np.mean(self.tarData, axis=3)
        self.noTarDataMean = np.mean(self.noTarData, axis=3)
        self.channelsNum = channelsNum
        self.personSignalLen = len(channelsNum)
        self.tarDataArrayIndex = 0
        self.nonTarDataArrayIndex = 0
        # self.tarDataIndex = 0
        # self.nonTarDataIndex = 0
        self.dataSize = self.dataTarArraySize + self.dataNonTargetArraySize
        self.data = np.empty((self.dataSize, 2 * self.personSignalLen, self.signalLen))
        self.label = np.empty(self.dataSize)
        count = 0
        while count < self.dataTarArraySize:
            subjectA = self.tarDataMean[self.tarDataArrayIndex]
            subjectB = self.tarDataMean[
                (self.tarDataArrayIndex + 1) % self.dataTarArraySize
            ]
            # if np.size(subjectA, 1) != np.size(subjectB, 1):
            #     minLen = min(np.size(subjectA, 1), np.size(subjectB, 1))
            # else:
            #     minLen = np.size(subjectA, 1)
            dataAList = []
            dataBList = []
            for i in self.channelsNum:
                dataAList.append(subjectA[i, 0 : self.signalLen])
            for i in self.channelsNum:
                dataBList.append(subjectB[i, 0 : self.signalLen])
            self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
            self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                np.array(dataBList)
            )
            self.label[count] = 1
            count += 1
            self.tarDataArrayIndex += 1
        while count - self.dataTarArraySize < self.dataNonTargetArraySize:
            subjectA = self.noTarDataMean[self.nonTarDataArrayIndex]
            subjectB = self.noTarDataMean[
                (self.nonTarDataArrayIndex + 1) % self.dataNonTargetArraySize
            ]
            # if np.size(subjectA, 1) != np.size(subjectB, 1):
            #     minLen = min(np.size(subjectA, 1), np.size(subjectB, 1))
            # else:
            #     minLen = np.size(subjectA, 1)
            dataAList = []
            dataBList = []
            for i in self.channelsNum:
                dataAList.append(subjectA[i, 0 : self.signalLen])
            for i in self.channelsNum:
                dataBList.append(subjectB[i, 0 : self.signalLen])
            self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
            self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                np.array(dataBList)
            )
            self.label[count] = 0
            count += 1
            self.nonTarDataArrayIndex += 1
        if count < self.dataSize:
            self.data = self.data[0:count, :, :]
            self.dataSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


class Dataset2022TarForMultiMean(Dataset):
    def __init__(self, channelsNum: Tuple, matPath: str) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        # self.signalLen = 307
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        self.tarData, self._ = dataset.getDataNdarray()
        self.signalLen = self.tarData.shape[2]
        # self.tarDataMean = np.empty((self.tarData.shape[0], self.tarData.shape[1], self.tarData.shape[2]))
        self.dataTarArraySize = self.tarData.shape[0]
        # for i in range(self.tarData.shape[0]):
        #     self.tarDataMean[i] = np.mean(self.tarData[i], axis=2)
        self.tarDataMean = np.mean(self.tarData, axis=3)
        self.channelsNum = channelsNum
        self.personSignalLen = len(channelsNum)
        self.tarDataArrayIndex = 0
        self.nonTarDataArrayIndex = 0
        # self.tarDataIndex = 0
        self.dataSize = self.dataTarArraySize
        self.data = np.empty((self.dataSize, 2 * self.personSignalLen, self.signalLen))
        self.label = np.empty(self.dataSize)
        count = 0
        while count < self.dataTarArraySize:
            subjectA = self.tarDataMean[self.tarDataArrayIndex]
            subjectB = self.tarDataMean[
                (self.tarDataArrayIndex + 1) % self.dataTarArraySize
            ]
            # if np.size(subjectA, 1) != np.size(subjectB, 1):
            #     minLen = min(np.size(subjectA, 1), np.size(subjectB, 1))
            # else:
            #     minLen = np.size(subjectA, 1)
            dataAList = []
            dataBList = []
            for i in self.channelsNum:
                dataAList.append(subjectA[i, 0 : self.signalLen])
            for i in self.channelsNum:
                dataBList.append(subjectB[i, 0 : self.signalLen])
            self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
            self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                np.array(dataBList)
            )
            self.label[count] = 1
            count += 1
            self.tarDataArrayIndex += 1

        if count < self.dataSize:
            self.data = self.data[0:count, :, :]
            self.dataSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


class Dataset2022NonTarForMultiMean(Dataset):
    def __init__(self, channelsNum: Tuple, matPath: str) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        # self.signalLen = 307
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        self._, self.noTarData = dataset.getDataNdarray()
        self.signalLen = self.noTarData.shape[2]
        # self.noTarDataMean = np.empty((self.noTarData.shape[0], self.noTarData.shape[1], self.noTarData.shape[2]))
        self.dataNonTargetArraySize = self.noTarData.shape[0]
        # for i in range(self.noTarData.shape[0]):
        #     self.noTarDataMean[i] = np.mean(self.noTarData[i], axis=2)
        self.noTarDataMean = np.mean(self.noTarData, axis=3)
        self.channelsNum = channelsNum
        self.personSignalLen = len(channelsNum)
        self.tarDataArrayIndex = 0
        self.nonTarDataArrayIndex = 0
        # self.tarDataIndex = 0
        # self.nonTarDataIndex = 0
        self.dataSize = self.dataNonTargetArraySize
        self.data = np.empty((self.dataSize, 2 * self.personSignalLen, self.signalLen))
        self.label = np.empty(self.dataSize)
        count = 0
        while count < self.dataNonTargetArraySize:
            subjectA = self.noTarDataMean[self.nonTarDataArrayIndex]
            subjectB = self.noTarDataMean[
                (self.nonTarDataArrayIndex + 1) % self.dataNonTargetArraySize
            ]
            # if np.size(subjectA, 1) != np.size(subjectB, 1):
            #     minLen = min(np.size(subjectA, 1), np.size(subjectB, 1))
            # else:
            #     minLen = np.size(subjectA, 1)
            dataAList = []
            dataBList = []
            for i in self.channelsNum:
                dataAList.append(subjectA[i, 0 : self.signalLen])
            for i in self.channelsNum:
                dataBList.append(subjectB[i, 0 : self.signalLen])
            self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
            self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                np.array(dataBList)
            )
            self.label[count] = 0
            count += 1
            self.nonTarDataArrayIndex += 1
        if count < self.dataSize:
            self.data = self.data[0:count, :, :]
            self.dataSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


class Dataset2022ForMultiSingle(Dataset):
    def __init__(
        self, channelsNum: Tuple, matPath: str
    ) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        # self.signalLen = 307
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        self.tarData, self.noTarData = dataset.getDataNdarray()
        self.signalLen = self.tarData.shape[2]
        self.dataTarSize, self.dataNonTargetSize = dataset.getDataSize()
        self.dataTarArraySize, self.dataNonTargetArraySize = dataset.getDataArraySize()
        self.channelsNum = channelsNum
        self.tarDataArrayIndex = 0
        self.nonTarDataArrayIndex = 0
        # self.tarDataIndex = 0
        # self.nonTarDataIndex = 0
        self.personSignalLen = len(channelsNum)
        self.dataSize = self.dataTarSize + self.dataNonTargetSize
        self.data = np.empty((self.dataSize, 2 * self.personSignalLen, self.signalLen))
        self.label = np.empty(self.dataSize)
        count = 0
        # print(self.tarData.shape)
        while count < self.dataTarSize:
            subjectA = self.tarData[self.tarDataArrayIndex]
            subjectB = self.tarData[self.tarDataArrayIndex]
            minLen = np.size(subjectA, 2)
            for i in range(minLen):
                dataAList = []  # 前3段信号
                dataBList = []  # 后3段信号
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
                for j in range(len(self.singleDataChannel)):
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
                self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
                self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                    np.array(dataBList)
                )
                self.label[count] = 1
                count += 1
            self.tarDataArrayIndex += 1
        while count - self.dataTarSize < self.dataNonTargetSize:
            subjectA = self.noTarData[self.nonTarDataArrayIndex]
            subjectB = self.noTarData[self.nonTarDataArrayIndex]
            minLen = np.size(subjectA, 2)
            for i in range(minLen):
                dataAList = []  # 前3段信号
                dataBList = []  # 后3段信号
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
                for j in range(len(self.singleDataChannel)):
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
                self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
                self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                    np.array(dataBList)
                )
                self.label[count] = 0
                count += 1
            self.nonTarDataArrayIndex += 1
        if count < self.dataSize:
            self.data = self.data[0:count, :, :]
            self.dataSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


class Dataset2022ForMultiWithSingle(Dataset):
    def __init__(
        self, channelsNum: Tuple, matPath: str, singlePersonFilePath: str
    ) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        # self.signalLen = 307
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        datasetSingle = Load2022Dataset.Load2022Dataset(singlePersonFilePath, True)
        self.tarData, self.noTarData = dataset.getDataNdarray()
        self.singleTarData, self.singleNoTarData = datasetSingle.getDataNdarray()
        self.signalLen = self.tarData.shape[2]
        self.dataTarSize, self.dataNonTargetSize = dataset.getDataSize()
        self.dataTarArraySize, self.dataNonTargetArraySize = dataset.getDataArraySize()
        self.channelsNum = channelsNum
        self.tarDataArrayIndex = 0
        self.nonTarDataArrayIndex = 0
        # self.tarDataIndex = 0
        # self.nonTarDataIndex = 0
        self.personSignalLen = len(channelsNum)
        self.dataSize = self.dataTarSize + self.dataNonTargetSize
        self.data = np.empty((self.dataSize, 2 * self.personSignalLen, self.signalLen))
        self.label = np.empty(self.dataSize)
        count = 0
        # print(self.tarData.shape)
        while count < self.dataTarSize:
            subjectA = self.tarData[self.tarDataArrayIndex]
            subjectB = self.singleTarData[0]
            if np.size(subjectA, 2) != np.size(subjectB, 2):
                minLen = min(np.size(subjectA, 2), np.size(subjectB, 2))
            else:
                minLen = np.size(subjectA, 2)
            choiceArray = [subjectA, subjectB]
            for i in range(minLen):
                dataAList = []
                dataBList = []
                curIndex = random.randint(0, 1)
                for j in self.channelsNum:
                    # dataAList.append(subjectA[j, 0:self.signalLen, i])
                    dataAList.append(
                        choiceArray[curIndex][j, 0 : self.signalLen, i]
                    )  # 单个人的信号随机选一个人的信号组成
                for j in self.channelsNum:
                    # dataBList.append(subjectB[j, 0:self.signalLen, i])
                    dataBList.append(
                        choiceArray[1 - curIndex][j, 0 : self.signalLen, i]
                    )
                self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
                self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                    np.array(dataBList)
                )
                self.label[count] = 1
                count += 1
            self.tarDataArrayIndex += 1
        while count - self.dataTarSize < self.dataNonTargetSize:
            subjectA = self.noTarData[self.nonTarDataArrayIndex]
            subjectB = self.singleNoTarData[0]
            choiceArray = [subjectA, subjectB]
            if np.size(subjectA, 2) != np.size(subjectB, 2):
                minLen = min(np.size(subjectA, 2), np.size(subjectB, 2))
            else:
                minLen = np.size(subjectA, 2)
            for i in range(minLen):
                dataAList = []
                dataBList = []
                curIndex = random.randint(0, 1)
                for j in self.channelsNum:
                    # dataAList.append(subjectA[j, 0:self.signalLen, i])
                    dataAList.append(
                        choiceArray[curIndex][j, 0 : self.signalLen, i]
                    )  # 单个人的信号随机选一个人的信号组成
                for j in self.channelsNum:
                    # dataBList.append(subjectB[j, 0:self.signalLen, i])
                    dataBList.append(
                        choiceArray[1 - curIndex][j, 0 : self.signalLen, i]
                    )
                self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
                self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = (
                    np.array(dataBList)
                )
                self.label[count] = 0
                count += 1
            self.nonTarDataArrayIndex += 1
        if count < self.dataSize:
            self.data = self.data[0:count, :, :]
            self.dataSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


class Dataset2022ForMultiWithMachine(Dataset):
    def __init__(
        self,
        channelsNum: Tuple,
        matPath: str,
        singlePersonFilePath: str,
        nchannel: list[int, int, int],
        labelCol: int,
        targetWord: str,
        csvDelimiter: str = "\t",
    ) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        # self.signalLen = 307
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        datasetSingle = LoadMachineData.LoadMachineDataSet(
            singlePersonFilePath, nchannel, labelCol, targetWord, csvDelimiter
        )
        self.tarData, self.noTarData = dataset.getDataNdarray()
        self.singleTarData, self.singleNoTarData = datasetSingle.getDataNdarray()
        self.singleDataChannel = datasetSingle.getSelectChoice()
        self.signalLen = self.tarData.shape[2]
        self.dataTarSize, self.dataNonTargetSize = dataset.getDataSize()
        self.channelsNum = channelsNum
        self.tarDataArrayIndex = 0
        self.nonTarDataArrayIndex = 0
        # self.tarDataIndex = 0
        # self.nonTarDataIndex = 0
        self.personSignalLen = len(channelsNum)
        self.dataSize = self.dataTarSize + self.dataNonTargetSize
        self.data = np.empty((self.dataSize, 2 * self.personSignalLen, self.signalLen))
        self.label = np.empty(self.dataSize)
        count = 0
        index = 0

        ta, _, _, td = self.tarData.shape
        # self.tarData=self.tarData.transpose(1,2,0,3).reshape(tb,tc,ta*td)

        subjectB = self.singleTarData
        targetDataLen = min(ta*td, np.size(subjectB, 2))
        for i in range(targetDataLen):
            if count != 0 and count % td == 0:
                index += 1
            subjectA = self.tarData[index]
            dataAList = []
            dataBList = []
            selectIndex = random.randint(0, 1)
            if selectIndex == 0:
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
                for j in range(len(self.singleDataChannel)):
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
            else:
                for j in range(len(self.singleDataChannel)):
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
            self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
            self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = np.array(dataBList)
            self.label[count] = 1
            count += 1

        index = 0
        nta, _, _, ntd = self.noTarData.shape
        # self.noTarData=self.noTarData.transpose(1,2,0,3).reshape(ntb,ntc,nta*ntd)
        nonTargetDataLen = min(nta * ntd, np.size(self.singleNoTarData, 2))
        subjectB = self.singleNoTarData
        for i in range(nonTargetDataLen):
            if count != 0 and count % ntd == 0:
                index += 1
            subjectA = self.noTarData[index]
            dataAList = []
            dataBList = []
            selectIndex = random.randint(0, 1)
            if selectIndex == 0:
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
                for j in range(len(self.singleDataChannel)):
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
            else:
                for j in range(len(self.singleDataChannel)):
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
            self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
            self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = np.array(dataBList)
            self.label[count] = 0
            count += 1

        if count < self.dataSize:
            self.data = self.data[0:count, :, :]
            self.label = self.label[0:count]
            self.dataSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


class Dataset2022ForMultiWithTargetMachine(Dataset):
    def __init__(
        self,
        channelsNum: Tuple,
        matPath: str,
        singlePersonFilePath: str,
        nchannel: list[int, int, int],
        labelCol: int,
        targetWord: str,
        csvDelimiter: str = "\t",
    ) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        # self.signalLen = 307
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        datasetSingle = LoadMachineData.LoadMachineDataSet(
            singlePersonFilePath, nchannel, labelCol, targetWord, csvDelimiter
        )
        self.tarData, _ = dataset.getDataNdarray()
        self.singleTarData, _ = datasetSingle.getDataNdarray()
        self.singleDataChannel = datasetSingle.getSelectChoice()
        self.signalLen = self.tarData.shape[2]
        self.dataTarSize, _ = dataset.getDataSize()
        self.channelsNum = channelsNum
        self.tarDataArrayIndex = 0
        self.personSignalLen = len(channelsNum)
        self.dataSize = self.dataTarSize
        self.data = np.empty((self.dataSize, 2 * self.personSignalLen, self.signalLen))
        self.label = np.empty(self.dataSize)
        count = 0
        # print(self.tarData.shape)
        a, _, _, d = self.tarData.shape

        subjectB = self.singleTarData
        targetDataLen = min(a * d, np.size(subjectB, 2))
        index = 0
        for i in range(targetDataLen):
            if count % d == 0 and count != 0:
                index += 1
            subjectA = self.tarData[index]
            dataAList = []
            dataBList = []
            selectIndex = random.randint(0, 1)
            if selectIndex == 0:
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
                for j in range(len(self.singleDataChannel)):
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
            else:
                for j in range(len(self.singleDataChannel)):
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
            self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
            self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = np.array(dataBList)
            self.label[count] = 1
            count += 1

        if count < self.dataSize:
            self.data = self.data[0:count, :, :]
            self.label = self.label[0:count]
            self.dataSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


class Dataset2022ForMultiWithNonTargetMachine(Dataset):
    def __init__(
        self,
        channelsNum: Tuple,
        matPath: str,
        singlePersonFilePath: str,
        nchannel: list[int, int, int],
        labelCol: int,
        targetWord: str,
        csvDelimiter: str = "\t",
    ) -> None:
        # [ch x time x trial]
        # self.signalLen = 610
        # self.signalLen = 307
        super().__init__()
        dataset = Load2022Dataset.Load2022Dataset(matPath, True)
        datasetSingle = LoadMachineData.LoadMachineDataSet(
            singlePersonFilePath, nchannel, labelCol, targetWord, csvDelimiter
        )
        _, self.noTarData = dataset.getDataNdarray()
        _, self.singleNoTarData = datasetSingle.getDataNdarray()
        self.singleDataChannel = datasetSingle.getSelectChoice()
        self.signalLen = self.noTarData.shape[2]
        _, self.dataNonTargetSize = dataset.getDataSize()
        self.channelsNum = channelsNum
        self.nonTarDataArrayIndex = 0
        self.personSignalLen = len(channelsNum)
        self.dataSize = self.dataNonTargetSize
        self.data = np.empty((self.dataSize, 2 * self.personSignalLen, self.signalLen))
        self.label = np.empty(self.dataSize)
        count = 0
        # print(self.tarData.shape)
        a, _, _, d = self.noTarData.shape
        nonTargetDataLen = min(a * d, np.size(subjectB, 2))
        subjectB = self.singleNoTarData
        for i in range(nonTargetDataLen):
            if count % d == 0 and count != 0:
                index += 1
            subjectA = self.noTarData[index]
            dataAList = []
            dataBList = []
            selectIndex = random.randint(0, 1)
            if selectIndex == 0:
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
                for j in range(len(self.singleDataChannel)):
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
            else:
                for j in range(len(self.singleDataChannel)):
                    dataBList.append(subjectB[j, 0 : self.signalLen, i])
                for j in self.channelsNum:
                    dataAList.append(subjectA[j, 0 : self.signalLen, i])
            self.data[count, 0 : self.personSignalLen, :] = np.array(dataAList)
            self.data[count, self.personSignalLen : 2 * self.personSignalLen, :] = np.array(dataBList)
            self.label[count] = 0
            count += 1

        if count < self.dataSize:
            self.data = self.data[0:count, :, :]
            self.label = self.label[0:count]
            self.dataSize = count
        for i in range(10):
            randomArray = np.random.permutation(count)
            self.data = self.data[randomArray]
            self.label = self.label[randomArray]

    def __len__(self) -> int:
        return self.dataSize

    def __getitem__(self, index: int) -> Tuple[np.ndarray, int]:
        return self.data[index], self.label[index]


if __name__ == "__main__":
    Dataset2022ForMulti((29, 30, 31), "/root/data/newData")
